#!/usr/bin/env python
import os
import sys
import re
import argparse
import shutil

class crysD3Gen():
    def __init__(self,title):
        self.title=title
        self.electrons='0'
        self.coreElectrons='0'
        self.totalbands=''
        self.shrink=1000
        self.nkpoints=''
        self.kpoint=''
        self.kpointLabel=''
        self.convergence=''
        self.dimension=3

    def readout(self,outFileName):
        outFile=open(outFileName,'r')
        shutil.copy(outFileName, outFileName+'.bak')
        for line in outFile:
            if 'SLAB' in line:
                self.dimension=2
            if 'POLYMER' in line:
                self.dimension=1
            if 'MOLECULE' in line:
                self.dimension=0
            if re.search('NUMBER OF AO',line):
                self.totalbands=line.split()[3]
            if re.search('N. OF ELECTRONS PER CELL', line):
                self.electrons=line.split()[5]
            if re.search('CORE ELECTRONS PER CELL', line):
                self.coreElectrons=line.split()[4]
            #if re.search('SHRINK. FACT.',line):
            #    self.shrink=map(int,[line.split()[2],line.split()[3],line.split()[4]])
            if re.search('NUMBER OF K POINTS IN THE IBZ', line):
                pos=line.split().index('IBZ')+1
                self.nkpoints=line.split()[pos]
            if re.search('\d-\w\(',line):
                label=re.sub('[^A-Z]','',line)
                self.kpointLabel=self.kpointLabel+label
                line=re.sub('\s+\d+-[A-Z]\(\s+','',line)
                line=re.sub('\)\s*',',',line)
                self.kpoint=self.kpoint+line    
            if re.search('SCF ENDED',line):
                self.convergence=line.split()[4]
                if self.convergence != 'CONVERGENCE':
                    print "Error: output file "+outFileName+" not ended normally"
                    sys.exit()           
        self.firstband=str(int(self.coreElectrons)/2+1)
        self.lastband=str(int(self.electrons)/2+10)
        self.nkpoints='100'
        self.nbands='1'
        if self.kpoint=='':
            print "Warning: not k point find in output file. Check the d3 file before running !!!"
        #print self.dimension
        kpoint1=str(self.shrink/2)+' 0 0'
        kpoint2='0 '+str(self.shrink/2)+' 0'
        if self.dimension==1:
            self.kpath='0 0 0\t'+kpoint1+'\n' 
        if self.dimension==2:
            self.kpath='0 0 0\t'+kpoint1+'\n' 
            self.kpath=self.kpath+kpoint1+'\t'+kpoint2+'\n' 
            self.kpath=self.kpath+kpoint2+'\t'+'0 0 0\n'
        self.dossShrink='16'
        self.ndosPT='100'
        self.ndoss='0'
      
                
    def writed3(self,d3FileName):
        d3File=open(d3FileName,'w')
        self.nbands=len(self.kpath.split('\n'))-1
        bandParam=str(self.nbands)+' '+str(self.shrink)+' '+self.nkpoints+' '+self.firstband+' '+self.lastband+' 1 1\n'
        d3File.write('BAND\n')
        d3File.write('input generated by QCflash\n')
        d3File.write(bandParam)
        d3File.write(self.kpath)
        d3File.write('NEWK\n')
        IS=int(self.dossShrink)
        d3File.write(str(IS)+' '+str(IS*2)+'\n')
        d3File.write('1 0\n')
        dossParam=self.ndoss+' '+self.ndosPT+' '+self.firstband+' '+self.lastband+' 2 20 0\n'
        d3File.write('DOSS\n')   
        d3File.write(dossParam)
        d3File.write('BWIDTH\n')
        d3File.write(self.firstband+' '+self.lastband+' \n')
        d3File.write('END')
        d3File.close()

parser=argparse.ArgumentParser(description='Generate d3 input file for crystal 14. Need to read the out file form scf calculaton (run .d12 file first)')
parser.add_argument("-r","--range", help="Two numbers to set the range for band and DOS. Default is all valence band to 10 about conductive band, e.g 10,5 means from 10 lower than HOMO to 5 higher than HOMO",type=str)
parser.add_argument("-K","--nkpoints",help='Number of k points along the k point path. Default: 100',type=str,default='100')
parser.add_argument("-D","--ndospoint",help='Number of points to calculate DOS. Default: 100',type=str,default='100')
parser.add_argument("-s","--shrink",help='Shrink for band structure. Default: 1000',type=int,default=1000)
parser.add_argument("-d","--dimension",help='Set the dimension for periodicity (0,1,2,3). Default: read from outfile', type=int)
parser.add_argument("-p","--kpath",help='k point path to draw in band structure. If more than one path seperate by semicolum. Defaule: 000 to last in out file. e.g.\"0 0 0;0.5 0 0;0 0 0.5\"')
parser.add_argument("-v","--version",action='version', version='%(prog)s 0.1')
parser.add_argument("outFiles",nargs='+')
args=parser.parse_args()
for outFile in args.outFiles:
    if not os.path.isfile(outFile):
        print outFile+" do not exist"
        sys.exit()
    outFileName=''.join(re.split('[/.]',outFile)[-2:-1])
    d3gen=crysD3Gen(outFileName)
    if args.dimension:
        d3gen.dimension=args.dimension
    d3gen.shrink=args.shrink
    d3gen.readout(outFileName+'.out')
    
    if args.nkpoints:
        d3gen.nkpoints=args.nkpoints
    if args.kpath:
        d3gen.kpath=''
        kpath=args.kpath.split(';')
        for kp in range(len(kpath)-1):
            d3gen.kpath=d3gen.kpath+' '.join([str(int(float(kpt)*d3gen.shrink))  for kpt in kpath[kp].split()])+'\t'+' '.join([str(int(float(kpt)*d3gen.shrink))  for kpt in kpath[kp+1].split()])+'\n'
    if args.ndospoint:
         d3gen.ndosPT=args.ndospoint
    if args.range:
        if str(int(d3gen.electrons)/2-int(args.range.split(',')[0]))>0:
            d3gen.firstband=str(int(d3gen.electrons)/2-int(args.range.split(',')[0]))
        else:
            d3gen.firstband='1'
            
        if str(int(d3gen.electrons)/2+int(args.range.split(',')[1]))<d3gen.totalbands:
            d3gen.lastband=str(int(d3gen.electrons)/2+int(args.range.split(',')[1]))
        else:
            d3gen.lastband=d3gen.totalbands
    print outFileName+": first band is "+str(d3gen.firstband)
    print outFileName+": last band is "+str(d3gen.lastband)
    d3gen.writed3(outFileName+'.d3')
